package com.katze.boot.applet.ssh;

import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.channel.ChannelExec;
import org.apache.sshd.client.channel.ChannelShell;
import org.apache.sshd.client.channel.ClientChannelEvent;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.common.util.security.SecurityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.security.GeneralSecurityException;
import java.time.Duration;
import java.util.EnumSet;
import java.util.function.BiConsumer;

/**
 * 连接Linux主机远程执行命令（同时支持交互与非交互两种形式）
 */
public class SSHExecutor implements Closeable {
    private static final Logger log = LoggerFactory.getLogger(SSHExecutor.class);

    private SshClient client;

    private ClientSession session;

    private String id;

    public static Builder builder() {
        return new Builder();
    }

    public static Builder builder(String host, String username, String password, Integer... port) {
        return new Builder(username, password, new InetSocketAddress(host, port.length > 0? port[0]:22));
    }

    /**
     * 密码方式
     * @param username 用户名
     * @param password 密码
     * @param address 主机地址
     */
    private SSHExecutor(String username, String password, InetSocketAddress address) {
        try {
            client = SshClient.setUpDefaultClient();
            client.start();
            session = client.connect(username, address).verify().getSession();
            session.addPasswordIdentity(password);
            id = address.getHostName() + "-" + address.getPort() + "-" + username;
        } catch (IOException e) {
            log.error("{}:{}:{}初始化异常", address, username, password, e);
        }
    }

    /**
     * 公钥方式
     * @param host
     * @param username
     * @param port
     * @param keyName
     * @param publicKey
     */
    private SSHExecutor(String host, String username, Integer port, String keyName, String publicKey) {
        try {
            client = SshClient.setUpDefaultClient();
            client.start();
            session = client.connect(username, host, port).verify().getSession();
            SecurityUtils.loadKeyPairIdentities(null, null, new ByteArrayInputStream(publicKey.getBytes()), null);
//            org.apache.sshd.common.util.security.SecurityUtils.loadKeyPairIdentity(keyName, new ByteArrayInputStream(publicKey.getBytes()), null)
            session.addPublicKeyIdentity(null);
        } catch (IOException | GeneralSecurityException e) {
            e.printStackTrace();
        }
    }

    /**
     * 用户登录认证
     * @param timeout 超时时长
     */
    protected SSHExecutor login(Duration timeout) throws IOException {
        session.auth().verify(timeout);
        return this;
    }

    /**
     * 非交互式执行命令
     * @param command 命令
     */
    public void execute(String command, BiConsumer<ByteArrayOutputStream, ByteArrayOutputStream> consumer) throws Exception{
        try (ByteArrayOutputStream out = new ByteArrayOutputStream(); ByteArrayOutputStream err = new ByteArrayOutputStream()){
            log.info("execute command: {}", command);
            ChannelExec channel = session.createExecChannel(command);
            channel.setOut(out);
            channel.setErr(err);
            channel.open();
            channel.waitFor(EnumSet.of(ClientChannelEvent.CLOSED), 0);
            channel.close(false); //true:立马关闭；false:优雅关闭
            consumer.accept(out, err);
        }
    }

    /**
     * 交互式执行命令
     */
    public void execute() {
        try {
            ChannelShell channel = session.createShellChannel();
            channel.setOut(System.out);
            channel.setErr(System.err);
            channel.setIn(System.in);
            channel.open();
            channel.waitFor(EnumSet.of(ClientChannelEvent.CLOSED), 0);
            channel.close(false);
        } catch (Exception e) {
            log.error("execute fail", e);
        }
    }

    @Override
    public void close (){
        session.close(false);
        client.stop();
    }

    public String getId() {
        return id;
    }

    public static class Builder {

        private SSHExecutor executor;

        private String username;

        private String password;

        private InetSocketAddress address;

        public Builder() {}

        public Builder(String username, String password, InetSocketAddress address) {
            this.username = username;
            this.password = password;
            this.address = address;
        }

        public Builder connection() {
            this.executor = new SSHExecutor(this.username, this.password, this.address);
            return this;
        }

        public Builder connection(String host, int port, String username, String password) {
            this.executor = new SSHExecutor(username, password, new InetSocketAddress(host, port));
            return this;
        }

        public void test() throws IOException {
            try {
                this.executor.login(Duration.ofSeconds(5));
            } finally {
                this.executor.close();
            }
        }

        public SSHExecutor build() throws IOException {
            return this.executor.login(Duration.ofSeconds(5));
        }

        public SSHExecutor build(Duration timeout) throws IOException {
            return this.executor.login(timeout);
        }
    }

    public static void main(String[] args) throws IOException {
        String host = "10.11.1.180",  username = "icstest", password = "ics#123.ITCO";
        try (SSHExecutor executor = SSHExecutor.builder(host, username, password).connection().build()){
            executor.execute();
        }
    }
}